package com.abc.config;

import com.abc.doc.KnowledgeDocumentPostProcessor;
import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatModel;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.rag.advisor.RetrievalAugmentationAdvisor;
import org.springframework.ai.rag.generation.augmentation.ContextualQueryAugmenter;
import org.springframework.ai.rag.preretrieval.query.transformation.RewriteQueryTransformer;
import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

@Configuration
public class RagConfig {

    @Autowired
    private VectorStore vectorStore;

    @Autowired
    private DashScopeChatModel dashScopeChatModel;

    @Autowired
    private KnowledgeDocumentPostProcessor knowledgeDocumentPostProcessor;


    @Bean
    public RetrievalAugmentationAdvisor retrievalAugmentationAdvisor() {

        return RetrievalAugmentationAdvisor.builder()
                .queryTransformers(RewriteQueryTransformer.builder()
                        .chatClientBuilder(ChatClient.builder(dashScopeChatModel).build().mutate())
                        .build())
                .documentRetriever(
                        VectorStoreDocumentRetriever.builder()
                                .similarityThreshold(0.30)
                                .vectorStore(vectorStore)
                                .build())
                .documentPostProcessors(knowledgeDocumentPostProcessor)
                .queryAugmenter(ContextualQueryAugmenter.builder()
                        .allowEmptyContext(false)
                        .build())
                .build();
    }
}
